from fastapi import FastAPI, Request, Form, HTTPException
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import BaseModel
import os
from pathlib import Path
from main import initialize_system, check_ollama_service
from typing import Dict, List, Optional
import logging
import json
import uuid
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建FastAPI应用
app = FastAPI(title="公共卫生知识库智能助理")

# 设置模板目录
templates = Jinja2Templates(directory="app/templates")

# 设置静态文件目录
app.mount("/static", StaticFiles(directory="app/static"), name="static")

# 初始化系统组件
system = None

# 对话历史存储
CHATS_FILE = "data/chats.json"
os.makedirs(os.path.dirname(CHATS_FILE), exist_ok=True)

# 定义请求模型
class ChatRequest(BaseModel):
    message: str
    history: Optional[List[Dict[str, str]]] = None
    chatId: Optional[str] = None

class ChatData(BaseModel):
    id: Optional[str] = None
    title: str
    messages: List[Dict[str, str]]
    created_at: Optional[str] = None
    updated_at: Optional[str] = None

# 加载对话历史
def load_chats():
    try:
        if os.path.exists(CHATS_FILE):
            with open(CHATS_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
        return []
    except Exception as e:
        logger.error(f"加载对话历史失败: {str(e)}")
        return []

# 保存对话历史
def save_chats(chats):
    try:
        with open(CHATS_FILE, 'w', encoding='utf-8') as f:
            json.dump(chats, f, ensure_ascii=False, indent=2)
    except Exception as e:
        logger.error(f"保存对话历史失败: {str(e)}")

# API路由
@app.get("/api/chats")
async def get_chats():
    """获取所有对话历史"""
    chats = load_chats()
    return chats

@app.get("/api/chats/{chat_id}")
async def get_chat(chat_id: str):
    """获取特定对话"""
    chats = load_chats()
    chat = next((c for c in chats if c['id'] == chat_id), None)
    if not chat:
        raise HTTPException(status_code=404, detail="对话不存在")
    return chat

@app.post("/api/chats")
async def create_chat(chat: ChatData):
    """创建新对话"""
    chats = load_chats()
    new_chat = chat.dict()
    new_chat['id'] = str(uuid.uuid4())
    new_chat['created_at'] = datetime.now().isoformat()
    new_chat['updated_at'] = new_chat['created_at']
    chats.insert(0, new_chat)
    save_chats(chats)
    return new_chat

@app.put("/api/chats")
async def update_chat(chat: ChatData):
    """更新现有对话"""
    if not chat.id:
        raise HTTPException(status_code=400, detail="缺少对话ID")
    
    chats = load_chats()
    chat_index = next((i for i, c in enumerate(chats) if c['id'] == chat.id), None)
    
    if chat_index is None:
        raise HTTPException(status_code=404, detail="对话不存在")
    
    updated_chat = chat.dict()
    updated_chat['updated_at'] = datetime.now().isoformat()
    chats[chat_index] = updated_chat
    save_chats(chats)
    return updated_chat

@app.delete("/api/chats/{chat_id}")
async def delete_chat(chat_id: str):
    """删除对话"""
    chats = load_chats()
    chats = [c for c in chats if c['id'] != chat_id]
    save_chats(chats)
    return {"status": "success"}

@app.on_event("startup")
async def startup_event():
    global system
    try:
        if check_ollama_service():
            logger.info("Ollama服务检查通过，开始初始化系统...")
            system = initialize_system()
            logger.info("系统初始化完成")
        else:
            logger.error("Ollama服务未运行")
            raise Exception("Ollama服务未运行，请确保Ollama服务已启动")
    except Exception as e:
        logger.error(f"系统初始化失败: {str(e)}")
        raise

@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    return templates.TemplateResponse(
        "index.html",
        {"request": request, "title": "公共卫生知识库智能助理"}
    )

@app.post("/ask")
async def ask_question(question: str = Form(...)):
    if not system:
        logger.error("系统未初始化")
        return JSONResponse(status_code=500, content={"error": "系统未初始化"})
    
    try:
        logger.info(f"处理问题: {question}")
        result = system['qa_system'].answer(question)
        logger.info("问题处理完成")
        return {
            "answer": result["answer"],
            "source": result["source"], 
            "references": result["references"]
        }
    except Exception as e:
        logger.error(f"处理问题时出错: {str(e)}")
        return JSONResponse(
            status_code=500, 
            content={"error": f"处理问题时出错: {str(e)}"}
        )

@app.post("/chat")
async def chat(request: ChatRequest):
    if not system:
        logger.error("系统未初始化")
        return JSONResponse(status_code=500, content={"error": "系统未初始化"})
    
    try:
        logger.info(f"处理聊天消息: {request.message}")
        if not request.history:
            request.history = []
        
        request.history.append({"role": "user", "content": request.message})
        result = system['qa_system'].chat(request.history)
        request.history.append({"role": "assistant", "content": result["answer"]})
        
        # 如果有对话ID，更新对话历史
        if request.chatId:
            chats = load_chats()
            chat_index = next((i for i, c in enumerate(chats) if c['id'] == request.chatId), None)
            if chat_index is not None:
                chats[chat_index]['messages'] = request.history
                chats[chat_index]['updated_at'] = datetime.now().isoformat()
                save_chats(chats)
        
        logger.info("聊天消息处理完成")
        return {
            "answer": result["answer"],
            "source": result["source"],
            "references": result["references"],
            "history": request.history
        }
    except Exception as e:
        logger.error(f"聊天时出错: {str(e)}")
        return JSONResponse(
            status_code=500, 
            content={"error": f"聊天时出错: {str(e)}"}
        )

@app.get("/health")
async def health_check():
    """
    健康检查接口，检查系统状态
    """
    if not system:
        return {"status": "error", "message": "系统未初始化"}
    
    return {
        "status": "ok", 
        "ollama_model": system['qa_system'].ollama_model_name,
        "vector_db_size": len(system['vector_store'].documents) if hasattr(system['vector_store'], 'documents') else 0
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)